import backtrader as bt
import datetime

class MyStrategy(bt.Strategy):
    params = (
        ("sma_period", 20),
    )

    def __init__(self):
        self.sma = bt.indicators.SimpleMovingAverage(self.data.close, period=self.params.sma_period)

    def next(self):
        print(f'self.data.close={self.data.close},self.sma={self.sma}')
        if self.data.close > self.sma:
            # 如果收盘价超过移动平均线，买入
            print(f'买入')
            self.buy()
        elif self.data.close < self.sma:
            # 如果收盘价低于移动平均线，卖出
            print(f'卖出')
            self.sell()

if __name__ == "__main__":
    cerebro = bt.Cerebro()

    # 加载数据
    data = bt.feeds.YahooFinanceCSVData(
        dataname='rb2401.csv',
        fromdate=datetime.datetime(2023, 11, 10),
        todate=datetime.datetime(2023, 11, 15),
        reverse=False
    )

    cerebro.adddata(data)

    # 添加策略
    cerebro.addstrategy(MyStrategy)

    # 设置初始资金
    cerebro.broker.set_cash(100000)

    # 设置手续费
    cerebro.broker.setcommission(commission=0.001)

    print(f"初始资金: {cerebro.broker.getvalue()}")




    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')

    result = cerebro.run()
    # 4.输出结果
    strats = result[0]

    print(f"结束资金: {cerebro.broker.getvalue()}")

    print('Sharpe Ratio:', strats.analyzers.sharpe.get_analysis())
    print('DrawDown:', strats.analyzers.drawdown.get_analysis())

    cerebro.plot(style='candlestick', barup='red', bardown='green', fmt_x_data='%Y-%m-%d', fmt_x_ticks='%Y-%m-%d')
